import struct
import random
from cryptography import x509
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import ec, padding
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

# 全局变量
g_server_hello_len = 0
g_certificate_len = 0
master_secret = 0
session_cache = {}


def create_server_responses(client_hello, cert_path="server.crt", key_path="server.key"):
    # 解析ClientHello获取必要参数（简化处理）
    client_version = client_hello[5:7]
    cipher_suite = struct.unpack('>H', client_hello[34:36])[0]
    session_id = client_hello[36 + 2 + 2 * int.from_bytes(client_hello[32:34], 'big') + 1: 36 + 2 + 2 * int.from_bytes(
        client_hello[32:34], 'big') + 1 + int.from_bytes(client_hello[36 + 2 + 2 * int.from_bytes(client_hello[32:34], 'big'):36 + 2 + 2 * int.from_bytes(
        client_hello[32:34], 'big') + 1], 'big')]

    # 检查会话恢复
    if session_id and session_id in session_cache:
        # 会话恢复
        print("Session resumption detected.")
        server_hello = create_server_hello(client_hello, session_id)
        server_hello_done = create_server_hello_done()
        responses = server_hello + server_hello_done
        master_secret = session_cache[session_id]
    else:
        # 新会话
        server_hello = create_server_hello(client_hello)
        global g_server_hello_len
        g_server_hello_len = len(server_hello) + 9 + 5
        # 生成Certificate
        certificate = create_certificate(cert_path)

        # 生成ServerKeyExchange（ECDHE专用）
        if cipher_suite in [0xC030, 0xC02C]:  # ECDHE-RSA-AES256-GCM-SHA384/ECDHE-ECDSA-AES256-GCM-SHA384
            server_key_exchange = create_server_key_exchange()
        else:
            server_key_exchange = create_server_key_exchange()

        # 生成ServerHelloDone
        server_hello_done = create_server_hello_done()

        # 合并所有响应
        responses = server_hello + certificate + server_key_exchange + server_hello_done

    return responses


def create_server_hello(client_hello, session_id=b''):
    client_version = client_hello[5:7]
    cipher_suites = client_hello[34:34 + 2 + 2 * int.from_bytes(client_hello[32:34], 'big')]

    record_layer = b'\x16\x03\x03'  # TLS 1.2 handshake message

    handshake = (
            b'\x02' +  # Handshake type: ServerHello
            b'\x00\x00\x00' +  # Handshake length (to be filled)
            client_version +  # Negotiated version
            struct.pack('>I', random.getrandbits(32)) +  # Server random timestamp
            bytes(random.getrandbits(8) for _ in range(28)) +  # Server random bytes
            struct.pack('>B', len(session_id)) +  # Session ID length
            session_id +  # Session ID
            struct.pack('>H', 0xC030) +  # Selected cipher suite (ECDHE-RSA-AES256-GCM-SHA384)
            b'\x00' +  # Compression method
            b'\x00\x00'  # Extensions (empty)
    )

    handshake_length = struct.pack('>I', len(handshake))[1:]  # Ignore first byte
    handshake = handshake[:1] + handshake_length + handshake[4:]

    record_length = struct.pack('>H', len(handshake))
    server_hello = record_layer + record_length + handshake

    return server_hello


def create_certificate(cert_path):
    # 读取证书文件
    with open(cert_path, 'rb') as f:
        cert_data = f.read()

    # 解析证书DER数据
    cert = x509.load_pem_x509_certificate(cert_data)
    der_bytes = cert.public_bytes(serialization.Encoding.DER)
    global g_certificate_len
    g_certificate_len = len(der_bytes)
    # 构造Certificate报文
    cert_length = struct.pack('>I', len(der_bytes))
    cert_count = struct.pack('>B', 1)
    handshake = (
            b'\x0B' +  # 握手类型：Certificate
            struct.pack('>I', len(cert_length) + len(cert_count) + len(der_bytes))[1:] +
            cert_length +
            cert_count +
            der_bytes
    )
    return wrap_handshake(handshake)


def create_server_key_exchange():
    # 生成ECDHE临时密钥对
    private_key = ec.generate_private_key(ec.SECP256R1())
    public_key = private_key.public_key()

    # 序列化公钥
    der_public = public_key.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )

    # 构造ServerKeyExchange
    curve_type = struct.pack('>H', 0x0017)  # secp256r1
    public_key_length = struct.pack('>H', len(der_public))
    handshake = (
            b'\x0C' +  # 握手类型：ServerKeyExchange
            struct.pack('>I', len(curve_type) + len(public_key_length) + len(der_public))[1:] +
            curve_type +
            public_key_length +
            der_public
    )
    return wrap_handshake(handshake)


def create_server_hello_done():
    handshake = (
            b'\x0E' +  # 握手类型：ServerHelloDone
            b'\x00\x00\x00'  # 握手长度（待填充）
    )
    handshake = handshake[:1] + struct.pack('>I', 0)[1:] + handshake[4:]
    return wrap_handshake(handshake)


def wrap_handshake(handshake_data):
    # 封装到TLS记录层
    record_type = b'\x16'
    protocol_version = b'\x03\x03'
    record_length = struct.pack('>H', len(handshake_data))
    return record_type + protocol_version + record_length + handshake_data


def create_client_responses(server_responses, client_hello, cert_path="server.crt"):
    # 解析服务器响应
    server_hello = server_responses[0: server_responses.find(b'\x0B')]
    certificate = server_responses[g_server_hello_len: g_server_hello_len + g_certificate_len]
    server_key_exchange = server_responses[server_responses.find(b'\x0C'): server_responses.find(b'\x0E')]
    server_hello_done = server_responses[server_responses.find(b'\x0E'):]

    # 解析服务器证书
    # cert_der = certificate[9: int.from_bytes(certificate[4:8], byteorder='big') + 9]
    cert = x509.load_der_x509_certificate(certificate)
    server_public_key = cert.public_key()

    # 生成ClientKeyExchange
    client_key_exchange = create_client_key_exchange(server_public_key)

    # 生成ChangeCipherSpec
    change_cipher_spec = create_change_cipher_spec()

    # 生成Finished
    finished = create_finished(client_hello, server_responses, client_key_exchange)

    return client_key_exchange + change_cipher_spec + finished


def create_client_key_exchange(server_public_key):
    # 生成ECDHE临时密钥对
    private_key = ec.generate_private_key(ec.SECP256R1())
    public_key = private_key.public_key()

    # 序列化公钥
    der_public = public_key.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )

    # 生成Pre-Master Secret（ECDHE专用）
    pre_master_secret = bytes(random.getrandbits(8) for _ in range(48))

    # 计算共享密钥
    shared_secret = private_key.exchange(ec.ECDH(), public_key)
    global master_secret
    # 派生Master Secret
    master_secret = HKDF(
        algorithm=hashes.SHA256(),
        length=48,
        salt=b'',
        info=b'tls12 master secret',
    ).derive(shared_secret)

    # 构造ClientKeyExchange
    handshake = (
        b'\x10' +  # 握手类型：ClientKeyExchange
        struct.pack('>I', len(der_public))[1:] +
        der_public
    )
    return wrap_handshake(handshake)


def create_change_cipher_spec():
    # 构造ChangeCipherSpec
    record = b'\x14\x03\x03\x00\x01\x01'
    return record


def create_finished(client_hello, server_responses, client_key_exchange):
    # 合并所有握手数据
    handshake_messages = client_hello + server_responses + client_key_exchange
    global master_secret
    # 计算Finished消息的verify_data
    verify_data = HKDF(
        algorithm=hashes.SHA256(),
        length=12,
        salt=b'',
        info=b'client finished',
    ).derive(master_secret)

    # 构造Finished
    handshake = (
            b'\x14' +  # 握手类型：Finished（实际应为0x14，但需注意类型冲突）
            struct.pack('>I', 12)[1:] +
            verify_data
    )
    return wrap_handshake(handshake)


# 生成 ServerChangeCipherSpec 消息
def create_server_change_cipher_spec():
    # 构造ChangeCipherSpec
    record = b'\x14\x03\x03\x00\x01\x01'
    return record


# 生成 ServerFinished 消息
def create_server_finished(client_hello, server_responses, client_key_exchange):
    # 合并所有握手数据
    handshake_messages = client_hello + server_responses + client_key_exchange
    global master_secret
    # 计算Finished消息的verify_data
    verify_data = HKDF(
        algorithm=hashes.SHA256(),
        length=12,
        salt=b'',
        info=b'server finished',
    ).derive(master_secret)

    # 构造Finished
    handshake = (
            b'\x14' +  # 握手类型：Finished
            struct.pack('>I', 12)[1:] +
            verify_data
    )
    return wrap_handshake(handshake)


# 派生加密密钥和 IV
def derive_keys(master_secret, client_random, server_random):
    def prf(secret, label, seed, length):
        p_hash = b''
        a = b''
        while len(p_hash) < length:
            a = hashes.Hash(hashes.SHA256(), backend=default_backend())
            a.update(a + label + seed)
            a = a.finalize()
            h = hashes.Hash(hashes.SHA256(), backend=default_backend())
            h.update(a + secret + seed)
            p_hash += h.finalize()
        return p_hash[:length]

    key_block = prf(master_secret, b"key expansion", client_random + server_random, 80)
    client_write_key = key_block[0:16]
    server_write_key = key_block[16:32]
    client_write_iv = key_block[32:48]
    server_write_iv = key_block[48:64]
    return client_write_key, server_write_key, client_write_iv, server_write_iv


# 加密数据
def encrypt_data(data, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    padded_data = data + b'\x00' * (16 - len(data) % 16)
    return encryptor.update(padded_data) + encryptor.finalize()


# 解密数据
def decrypt_data(data, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    decrypted_data = decryptor.update(data) + decryptor.finalize()
    return decrypted_data.rstrip(b'\x00')


# 生成 ClientHello 报文
def create_client_hello(hostname="example.com", session_id=b''):
    record_type = b'\x16'  # Handshake
    protocol_version = b'\x03\x03'  # TLS 1.2
    length = b'\x00\x00'  # Will be filled later

    handshake_type = b'\x01'  # ClientHello
    handshake_length = b'\x00\x00\x00'  # Will be filled later
    client_version = b'\x03\x03'

    random_bytes = struct.pack('>I', random.getrandbits(32))  # GMT Unix time
    random_bytes += bytes(random.getrandbits(8) for _ in range(28))  # Random bytes

    cipher_suites = [
        0xC0, 0x30,  # ECDHE-RSA-AES256-GCM-SHA384
        0xC0, 0x2F,  # ECDHE-RSA-AES128-GCM-SHA256
        0x00, 0x9C  # AES256-GCM-SHA384
    ]
    cipher_suites_length = struct.pack('>H', len(cipher_suites))

    compression_methods = [0x00]  # NULL compression
    compression_methods_length = struct.pack('>B', len(compression_methods))

    extensions = []

    sni_hostname = hostname.encode('ascii')
    sni_extension = [
        b'\x00\x00',  # Type: server_name (0)
        struct.pack('>H', 5 + len(sni_hostname)),  # Length
        b'\x00\x00',  # Server name list length
        struct.pack('>H', len(sni_hostname)),  # Hostname length
        sni_hostname  # Hostname bytes
    ]
    extensions.extend(sni_extension)

    elliptic_curves = [
        b'\x00\x0A',  # Type: elliptic_curves (10)
        struct.pack('>H', 4),  # Length
        struct.pack('>H', 0x0017),  # secp256r1
        struct.pack('>H', 0x0018)  # secp384r1
    ]
    extensions.extend(elliptic_curves)

    extensions_length = struct.pack('>H', len(extensions))
    extensions_bytes = extensions_length + b''.join(extensions)

    client_hello = (
            handshake_type +
            handshake_length +
            client_version +
            random_bytes +
            struct.pack('>B', len(session_id)) +
            session_id +
            cipher_suites_length +
            bytes(cipher_suites) +
            compression_methods_length +
            bytes(compression_methods) +
            extensions_bytes
    )

    handshake_length = struct.pack('>I', len(client_hello))[1:]  # Skip first byte
    client_hello = client_hello[:1] + handshake_length + client_hello[4:]

    record_length = struct.pack('>H', len(client_hello))
    record_layer = record_type + protocol_version + record_length + client_hello

    return record_layer

# 示例使用
if __name__ == "__main__":
    client_hello = create_client_hello("example.com")
    server_responses = create_server_responses(client_hello)
    client_responses = create_client_responses(server_responses, client_hello)

    server_change_cipher_spec = create_server_change_cipher_spec()
    server_finished = create_server_finished(client_hello, server_responses, client_responses[:client_responses.find(b'\x14\x03\x03')])

    full_server_responses = server_responses + server_change_cipher_spec + server_finished

    print("Full Server Responses Hexdump:")
    print(' '.join(f'{b:02X}' for b in full_server_responses))